import os

class MemoryLib:
    def __init__(self):
        # 初始化时检查并创建 memory 文件夹
        self.memory_dir = os.path.join(os.getcwd(), 'memory')
        os.makedirs(self.memory_dir, exist_ok=True)
        self.toolDescription = [
            {
                "type":"function",
                "function": {
                    "name":"get_memory_list",
                    "description": "检索知识库里面的所有知识关键词列表的方法。如果你想知道知识库里面有哪些知识就可以用这个方法。",
                    "parameters": {
                        "type": "object",
                        "properties": {},
                        "required": []
                    }
                }
            },
            {
                "type":"function",
                "function": {
                    "name":"read_memory",
                    "description": "读取知识库中对应知识的方法。需要你指定想要知道的知识的关键词。",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "topic": {
                                "type": "string",
                                "description": "想要获取的知识库中知识的关键词"
                            }
                        },
                        "required": ["topic"],
                        "additionalProperties": False
                    },
                    "strict": True
                }
            },
            {
                "type":"function",
                "function": {
                    "name":"write_memory",
                    "description": "写入知识的方法。需要指定知识的关键词和知识的内容",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "topic": {
                                "type": "string",
                                "description": "想要写入的知识的关键词"
                            },
                            "content": {
                                "type": "string",
                                "description": "想要写入的知识的内容"
                            }
                        },
                        "required": ["topic", "content"],
                        "additionalProperties": False
                    },
                    "strict": True
                }
            }
        ]
    
    def get_memory_list(self):
        # 获取 memory 目录下所有文件的列表
        files = []
        for item in os.listdir(self.memory_dir):
            item_path = os.path.join(self.memory_dir, item)
            if os.path.isfile(item_path):
                files.append(item)
        return "知识库中的记录有: " + str(files)
    
    def read_memory(self, topic):
        # 读取指定文件内容
        file_path = os.path.join(self.memory_dir, topic)
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.read()
        return topic + "的内容是: " +content
    
    def write_memory(self, topic, content):
        # 写入/覆盖文件内容
        file_path = os.path.join(self.memory_dir, topic)
        with open(file_path, 'w', encoding='utf-8') as f:
            f.write(content)
        return topic + "已经记录到知识库。"
    
    def call_tool(self, name, args):
        # 先找到对应的方法名
        if name == 'get_memory_list':
            return self.get_memory_list()
        if name == 'read_memory':
            return self.read_memory(**args)
        if name == 'write_memory':
            return self.write_memory(**args)
        
    def get_character_setting(self):
        file_path = os.path.join(os.getcwd(), 'resource/characterSetting')
        with open(file_path, 'r', encoding='utf-8') as f:
            return f.read()